# -*- coding: utf-8 -*-
"""
    Project name：code_potentialflow
# -------------------------------
    File name：SCHALFPLANE.py
    Created on：2025/6/5 16:38
    Author：(input)
    Description:
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.tri as tri



class BodyFittedGrid:
    """Class for generating body-fitted grid using bilinear interpolation."""
    def __init__(self, L=10, h=2, alpha_deg=80, m=50, n=50):
        """
        Initialize the body-fitted grid generator.

        Parameters:
        - L: Length of the physical domain (x-coordinate of right boundary)
        - h: Height of the physical domain (y-coordinate of top boundary)
        - alpha_deg: Tilt angle of the left boundary (degrees)
        - m: Number of points in xi direction (vertical)
        - n: Number of points in eta direction (horizontal)
        """
        self.L = L
        self.h = h
        self.alpha = np.deg2rad(alpha_deg)
        self.m = m
        self.n = n
        self.X = None
        self.Y = None
        self.W = None

        # Define physical domain vertices
        self.P = np.array([
            [L, h],  # A1: Top-right (xi=1, eta=1)
            [-h * np.tan(self.alpha), h],  # A2: Top-left (xi=0, eta=1)
            [0.0, 0.0],  # A3: Bottom-left (xi=0, eta=0)
            [L, 0.0]  # A4: Bottom-right (xi=1, eta=0)
        ])

    def shape_functions(self, xi, eta):
        """Compute bilinear shape functions, aligned with vertices A1, A2, A3, A4."""
        N0 = xi * eta  # A1 (xi=1, eta=1)
        N1 = (1 - xi) * eta  # A2 (xi=0, eta=1)
        N2 = (1 - xi) * (1 - eta)  # A3 (xi=0, eta=0)
        N3 = xi * (1 - eta)  # A4 (xi=1, eta=0)
        return np.array([N0, N1, N2, N3])

    def map_to_physical(self, xi, eta):
        """Map computational coordinates (xi, eta) to physical coordinates (x, y)."""
        N = self.shape_functions(xi, eta)
        x = np.dot(N, self.P[:, 0])
        y = np.dot(N, self.P[:, 1])
        return x, y

    def map_to_upper_half(self, xi, eta):
        """Map computational coordinates (xi, eta) to upper half-plane (Re(w), Im(w))."""
        x, y = self.map_to_physical(xi, eta)
        z = x + 1j * y
        w = np.cosh(np.pi / self.h * z)
        return w.real, w.imag

    def generate_grid(self):
        """Generate body-fitted grid coordinates."""
        xi = np.linspace(0, 1, self.m)
        eta = np.linspace(0, 1, self.n)
        Xi, Eta = np.meshgrid(xi, eta, indexing='ij')

        self.X = np.zeros((self.m, self.n))
        self.Y = np.zeros((self.m, self.n))

        for i in range(self.m):
            for j in range(self.n):
                self.X[i, j], self.Y[i, j] = self.map_to_physical(Xi[i, j], Eta[i, j])

        return self.X, self.Y

    def map_to_upper_half_grid(self):
        """Map the entire grid to the upper half-plane."""
        if self.X is None or self.Y is None:
            self.generate_grid()

        z = self.X + 1j * self.Y
        A = np.pi / (self.h * (1+1j*np.tan(self.alpha)))
        self.W = np.cosh(A * z)
        return self.W

    def check_grid_quality(self):
        """Check grid quality for degeneration or severe distortion."""
        if self.X is None or self.Y is None:
            self.generate_grid()

        dx_dxi = np.diff(self.X, axis=0)
        dy_dxi = np.diff(self.Y, axis=0)
        dx_deta = np.diff(self.X, axis=1)
        dy_deta = np.diff(self.Y, axis=1)

        min_dim = min(dx_dxi.shape[0], dx_deta.shape[1])
        dx_dxi = dx_dxi[:min_dim, :min_dim]
        dy_dxi = dy_dxi[:min_dim, :min_dim]
        dx_deta = dx_deta[:min_dim, :min_dim]
        dy_deta = dy_deta[:min_dim, :min_dim]

        jacobian = dx_dxi * dy_deta - dy_dxi * dx_deta
        min_jacobian = np.min(np.abs(jacobian))

        if min_jacobian < 1e-6:
            print("Warning: Grid may be degenerate, minimum Jacobian determinant near zero")
        else:
            print(f"Grid quality normal, minimum Jacobian determinant: {min_jacobian:.6f}")

    def visualize(self, figsize=(8, 4), title="Body-Fitted Grid"):
        """Visualize the original grid with correct boundary labels."""
        if self.X is None or self.Y is None:
            self.generate_grid()

        plt.figure(figsize=figsize)

        # Plot grid lines
        for i in range(self.X.shape[0]):
            plt.plot(self.X[i, :], self.Y[i, :], color='gray', linewidth=0.5)
        for j in range(self.X.shape[1]):
            plt.plot(self.X[:, j], self.Y[:, j], color='gray', linewidth=0.5)

        # Plot boundaries with correct labels
        plt.plot(self.X[:, 0], self.Y[:, 0], 'r', label='Bottom')  # eta=0, A3 to A4
        plt.plot(self.X[:, -1], self.Y[:, -1], 'b', label='Top')  # eta=1, A2 to A1
        plt.plot(self.X[0, :], self.Y[0, :], 'g', label='Left')  # xi=0, A3 to A2
        plt.plot(self.X[-1, :], self.Y[-1, :], 'm', label='Right')  # xi=1, A4 to A1

        plt.legend()
        plt.axis('equal')
        plt.title(f"{title} (α={np.degrees(self.alpha):.1f}°)")
        plt.xlabel("x")
        plt.ylabel("y")
        plt.grid(True)
        plt.show()

    def visualize_mapped(self, figsize=(8, 4), title="Upper Half-Plane Grid"):
        """Visualize the grid mapped to the upper half-plane with correct boundary labels."""
        if self.W is None:
            self.map_to_upper_half_grid()

        W_real = self.W.real
        W_imag = self.W.imag

        plt.figure(figsize=figsize)

        # Plot grid lines
        for i in range(self.W.shape[0]):
            plt.plot(W_real[i, :], W_imag[i, :], color='gray', linewidth=0.5)
        for j in range(self.W.shape[1]):
            plt.plot(W_real[:, j], W_imag[:, j], color='gray', linewidth=0.5)

        # Plot boundaries with correct labels
        plt.plot(W_real[:, 0], W_imag[:, 0], 'r', label='Bottom')  # eta=0, A3 to A4
        plt.plot(W_real[:, -1], W_imag[:, -1], 'b', label='Top')  # eta=1, A2 to A1
        plt.plot(W_real[0, :], W_imag[0, :], 'g', label='Left')  # xi=0, A3 to A2
        plt.plot(W_real[-1, :], W_imag[-1, :], 'm', label='Right')  # xi=1, A4 to A1

        plt.legend()
        plt.axis('equal')
        plt.title(f"{title} (α={np.degrees(self.alpha):.1f}°)")
        plt.xlabel("Re(w)")
        plt.ylabel("Im(w)")
        plt.grid(True)
        plt.show()

    def get_boundary_points(self, N):
        """Extract points on three boundaries (top, left, bottom) and print vertex coordinates in physical and upper half-planes."""
        xi = np.linspace(0, 1, N)
        eta = np.linspace(0, 1, N)

        # Top boundary: eta=1, from A2 to A1
        top_points = np.array([self.map_to_physical(x, 1.0) for x in xi])
        # Left boundary: xi=0, from A3 to A2
        left_points = np.array([self.map_to_physical(0.0, e) for e in eta])
        # Bottom boundary: eta=0, from A3 to A4
        bottom_points = np.array([self.map_to_physical(x, 0.0) for x in xi])

        # Vertex mapping verification: physical and upper half-planes
        corners = {
            'A1 (L, h)': {'xi': 1, 'eta': 1, 'index': 0},  # Top-right
            'A2 (-h*tanα, h)': {'xi': 0, 'eta': 1, 'index': 1},  # Top-left
            'A3 (0, 0)': {'xi': 0, 'eta': 0, 'index': 2},  # Bottom-left
            'A4 (L, 0)': {'xi': 1, 'eta': 0, 'index': 3}  # Bottom-right
        }
        print("\n--- Vertex Coordinates Comparison (Physical and Upper Half-Plane) ---")
        for name, info in corners.items():
            xi_val, eta_val, idx = info['xi'], info['eta'], info['index']

            # Physical plane coordinates
            phys_mapped = self.map_to_physical(xi_val, eta_val)
            phys_defined = self.P[idx]
            phys_error = np.array(phys_mapped) - phys_defined

            # Upper half-plane coordinates
            upper_mapped = self.map_to_upper_half(xi_val, eta_val)

            # print(f"{name}:")
            # print(f"  Physical Plane Defined = {phys_defined}")
            # print(f"  Physical Plane Mapped = {np.array(phys_mapped)}")
            # print(f"  Physical Plane Error = {phys_error}")
            # print(f"  Upper Half-Plane Mapped = {np.array(upper_mapped)}\n")

        return {
            'top': top_points,
            'left': left_points,
            'bottom': bottom_points
        }


class ComplexPotential:
    def __init__(self, mapper, source_phys_pos, Qc, margin=0.1, search_N=100):
        """
        Parameters:
        - mapper: Must provide map_to_physical(xi, eta) and map_to_upper_half(xi, eta) methods
        - source_phys_pos: Source position in the physical plane (x, y)
        - Qc: Source strength
        - margin: Plot margin ratio
        - search_N: Search precision for source point inverse mapping
        """
        self.mapper = mapper
        self.source_phys = np.array(source_phys_pos)
        self.Qc = Qc
        self.margin = margin
        self.search_N = search_N
        self.h = self.mapper.h  # Use mapper's physical height
        self.compute_source_position()

    def compute_source_position(self):
        """
        Search for (xi, eta) in the computational domain closest to the physical source point,
        and compute the corresponding w_c in the upper half-plane.
        """
        x0, y0 = self.source_phys
        xi_vals = np.linspace(0, 1, self.search_N)
        eta_vals = np.linspace(0, 1, self.search_N)

        min_dist = float("inf")
        best_coords = (0.0, 0.0)
        for xi in xi_vals:
            for eta in eta_vals:
                x1, y1 = self.mapper.map_to_physical(xi, eta)
                dist = np.hypot(x1 - x0, y1 - y0)
                if dist < min_dist:
                    min_dist = dist
                    best_coords = (xi, eta)

        self.xi_eta_source = best_coords
        w_real, w_imag = self.mapper.map_to_upper_half(*best_coords)
        self.wc = complex(w_real, w_imag)

    def map_to_w(self, xi, eta):
        """
        Map computational coordinates (xi, eta) to upper half-plane w (complex).
        Alias for map_to_upper_half, returning complex number.
        """
        w_real, w_imag = self.mapper.map_to_upper_half(xi, eta)
        return complex(w_real, w_imag)

    def compute_field(self, Nx=100, Ny=100):
        """
        Generate velocity potential, stream function, and velocity components over the computational domain.
        """
        xi_vals = np.linspace(0, 1, Nx)
        eta_vals = np.linspace(0, 1, Ny)

        data = []
        for eta in eta_vals:
            for xi in xi_vals:
                x_phys, y_phys = self.mapper.map_to_physical(xi, eta)
                w = self.map_to_w(xi, eta)

                W = (self.Qc / np.pi) * np.log(w - self.wc)
                dW_dw = (self.Qc / np.pi) / (w - self.wc)

                phi = np.real(W)
                psi = np.imag(W)
                u = np.real(dW_dw)
                v = np.imag(dW_dw)
                V = np.hypot(u, v)

                data.append({
                    "x": x_phys,
                    "y": y_phys,
                    "phi": phi,
                    "psi": psi,
                    "u": u,
                    "v": v,
                    "V": V
                })

        df = pd.DataFrame(data)
        return df

    def plot_streamlines(self, df, levels=30, h=None, L=None, alpha_deg=45, boundary_N=200):
        """
        Plot streamlines with physical boundaries and source point.
        """
        x, y, psi = df["x"], df["y"], df["psi"]
        margin_x = self.margin * (x.max() - x.min())
        margin_y = self.margin * (y.max() - y.min())

        triang = tri.Triangulation(x, y)
        plt.figure(figsize=(7, 6))

        # Plot streamlines
        cs = plt.tricontour(triang, psi, levels=levels, colors='k', linewidths=0.6)
        plt.clabel(cs, inline=True, fontsize=8, fmt="%.1f")

        # Plot boundaries
        boundaries = self.mapper.get_boundary_points(boundary_N)
        plt.plot(boundaries['top'][:, 0], boundaries['top'][:, 1], 'k-', linewidth=2.5, label='Top wall')
        plt.plot(boundaries['left'][:, 0], boundaries['left'][:, 1], 'k-', linewidth=2.5, label='Inclined wall')
        plt.plot(boundaries['bottom'][:, 0], boundaries['bottom'][:, 1], 'k-', linewidth=2.5, label='Bottom wall')

        # Plot x=0 reference line
        y_ref = np.linspace(y.min(), y.max(), 200)
        x_ref = np.zeros_like(y_ref)
        plt.plot(x_ref, y_ref, 'b--', linewidth=1.8, label='x = 0 ref line')

        # Plot source point
        source_x, source_y = self.source_phys
        plt.plot(source_x, source_y, 'ro', markersize=6, label='Source Point')

        # Plot settings
        plt.title("Streamlines (ψ) with Physical Boundaries")
        plt.xlim(x.min() - margin_x, x.max() + margin_x)
        plt.ylim(y.min() - margin_y, y.max() + margin_y)
        plt.gca().set_aspect('equal')
        plt.xlabel("x")
        plt.ylabel("y")
        plt.grid(True)
        plt.legend(loc="best")
        plt.tight_layout()
        plt.show()

    def plot_streamlines_plain(self, df, levels=30):
        """
        Plot streamlines without boundaries or annotations.
        """
        x, y, psi = df["x"], df["y"], df["psi"]
        margin_x = self.margin * (x.max() - x.min())
        margin_y = self.margin * (y.max() - y.min())

        triang = tri.Triangulation(x, y)
        plt.figure(figsize=(6, 5))
        plt.tricontour(triang, psi, levels=levels, colors='k', linewidths=0.6)
        plt.title("Streamlines (ψ)")
        plt.xlim(x.min() - margin_x, x.max() + margin_x)
        plt.ylim(y.min() - margin_y, y.max() + margin_y)
        plt.gca().set_aspect('equal')
        plt.xlabel("x")
        plt.ylabel("y")
        plt.grid(True)
        plt.tight_layout()
        plt.show()

    def plot_velocity_vectors(self, df, stride=5, scale=1.0):
        """
        Plot velocity vector field.
        """
        x = df["x"].values.reshape(-1)
        y = df["y"].values.reshape(-1)
        u = df["u"].values.reshape(-1)
        v = df["v"].values.reshape(-1)

        plt.figure(figsize=(6, 5))
        plt.quiver(x[::stride], y[::stride], u[::stride], v[::stride], scale=scale)
        plt.title("Velocity Vectors")
        plt.xlabel("x")
        plt.ylabel("y")
        plt.gca().set_aspect('equal')
        plt.grid(True)
        plt.tight_layout()
        plt.show()

    def save_to_csv(self, df, filename="field_output.csv"):
        """
        Save flow field data to CSV.
        """
        df.to_csv(filename, index=False)
        print(f"[Saved] CSV file: {filename}")

    def extract_line_profile(self, df, fixed_coord='x', fixed_value=0.5, var='V', num=200):
        """
        Extract profile of a variable along a line with fixed x or y in the physical plane.
        Parameters:
        - df: Flow field DataFrame
        - fixed_coord: 'x' or 'y' (fixed coordinate)
        - fixed_value: Fixed value (e.g., x=0.5)
        - var: Variable to extract (e.g., 'u', 'v', 'V', 'phi', 'psi')
        - num: Number of interpolation points
        Returns:
        - (coords, values): Variable coordinates and values
        """
        if fixed_coord not in ['x', 'y']:
            raise ValueError("fixed_coord must be 'x' or 'y'")

        tol = 1e-3  # Tolerance
        filtered = df[np.abs(df[fixed_coord] - fixed_value) < tol]

        if fixed_coord == 'x':
            variable_coord = 'y'
        else:
            variable_coord = 'x'

        sorted_data = filtered.sort_values(by=variable_coord)
        coords = sorted_data[variable_coord].values
        values = sorted_data[var].values

        # Normalize values
        values = values / np.max(np.abs(values)) if np.max(np.abs(values)) > 1e-6 else values

        return coords, values

    def plot_profile_along_line(self, df, fixed_coord='x', fixed_value=0.5, var='V', label=None):
        """
        Plot profile of a variable along a line.
        """
        coords, values = self.extract_line_profile(df, fixed_coord, fixed_value, var)

        plt.figure(figsize=(6, 4))
        plt.plot(coords, values, 'b-', lw=2, label=label or f"{var} profile")
        plt.xlabel("y" if fixed_coord == 'x' else "x")
        plt.ylabel(f"{var} / max({var})")
        plt.title(f"Profile of {var} along {fixed_coord} = {fixed_value}")
        plt.grid(True)
        plt.tight_layout()
        if label:
            plt.legend()
        plt.show()

# Example usage
if __name__ == "__main__":
    # Test with alpha = 10 degrees
    length, height = 10, 3
    alpha_deg = 10

    U0 = -4.0
    b = 0.02
    x_c = 0.3
    Q_strength = U0 * b
    print(f"\nTest with alpha = {alpha_deg} degrees")
    try:
        grid = BodyFittedGrid(L=length, h=height, alpha_deg=alpha_deg, m=50, n=50)
        print("Physical plane vertices:", grid.P)
        grid.generate_grid()
        grid.check_grid_quality()
        #grid.visualize(title="Original Grid")
        grid.map_to_upper_half_grid()
        #grid.visualize_mapped(title="Upper Half-Plane Grid")

        potential = ComplexPotential(grid, source_phys_pos=(x_c, height), Qc=Q_strength, search_N=100)
        print("Source (xi, eta):", potential.xi_eta_source)
        print("Upper half-plane w_c:", potential.wc)

        # Compute flow field
        df = potential.compute_field(Nx=50, Ny=50)

        # Plot streamlines with boundaries
        potential.plot_streamlines(df, levels=20, boundary_N=100)

        # Plot plain streamlines
        potential.plot_streamlines_plain(df, levels=20)

        # Plot velocity vectors
        potential.plot_velocity_vectors(df, stride=3, scale=10.0)

        # Save data to CSV
        potential.save_to_csv(df, filename="flow_field.csv")

        # Plot profile along x=5
        potential.plot_profile_along_line(df, fixed_coord='x', fixed_value=5.0, var='V', label='Velocity')

        boundary_points = grid.get_boundary_points(N=50)
        plt.figure(figsize=(8, 4))
        for key, points in boundary_points.items():
            plt.plot(points[:, 0], points[:, 1], label=key, linewidth=2)
        # plt.legend()
        # plt.axis('equal')
        # plt.title(f"Boundaries (α={alpha_deg}°)")
        # plt.xlabel("x")
        # plt.ylabel("y")
        # plt.grid(True)
        # plt.show()
    except ValueError as e:
        print(f"Error: {e}")